import type { Base } from "../Base";
import { mat4 } from "wgpu-matrix";
import distanceSizedPointsVertWGSL from "./distance-sized-points.vert.wgsl";
import fixedSizePointsVertWGSL from "./fixed-sized-points.vert.wgsl";
import orangeFragWGSL from "./orange.frag.wgsl";
import texturedFragWGSL from "./textured.frag.wgsl";
import { quitIfWebGPUNotAvailable } from "../../utils";

// See: https://www.google.com/search?q=fibonacci+sphere
function createFibonacciSphereVertices({
  numSamples,
  radius,
}: {
  numSamples: number;
  radius: number;
}) {
  const vertices = [];
  const increment = Math.PI * (3 - Math.sqrt(5));
  for (let i = 0; i < numSamples; ++i) {
    const offset = 2 / numSamples;
    const y = i * offset - 1 + offset / 2;
    const r = Math.sqrt(1 - Math.pow(y, 2));
    const phi = (i % numSamples) * increment;
    const x = Math.cos(phi) * r;
    const z = Math.sin(phi) * r;
    vertices.push(x * radius, y * radius, z * radius);
  }
  return new Float32Array(vertices);
}
export default class Points implements Base {
  name = "Points";
  description?: string | undefined;
  settings = {
    fixedSize: false,
    textured: false,
    size: 10,
  };
  async run(canvas?: HTMLCanvasElement) {
    const self = this;
    const adapter = await navigator.gpu?.requestAdapter({
      featureLevel: "compatibility",
    });
    const device = await adapter?.requestDevice();
    quitIfWebGPUNotAvailable(adapter, device ?? null);

    // Get a WebGPU context from the canvas and configure it
    const context = canvas!.getContext("webgpu");
    const devicePixelRatio = window.devicePixelRatio;
    canvas!.width = canvas!.clientWidth * devicePixelRatio;
    canvas!.height = canvas!.clientHeight * devicePixelRatio;
    const presentationFormat = navigator.gpu.getPreferredCanvasFormat();
    context!.configure({
      device: device!,
      format: presentationFormat,
    });

    // Create a bind group layout so we can share the bind groups
    // with multiple pipelines.
    const bindGroupLayout = device!.createBindGroupLayout({
      entries: [
        {
          binding: 0,
          visibility: GPUShaderStage.VERTEX,
          buffer: {},
        },
        {
          binding: 1,
          visibility: GPUShaderStage.FRAGMENT,
          sampler: {},
        },
        {
          binding: 2,
          visibility: GPUShaderStage.FRAGMENT,
          texture: {},
        },
      ],
    });

    const pipelineLayout = device!.createPipelineLayout({
      bindGroupLayouts: [bindGroupLayout],
    });

    // Compile all 4 shaders
    const shaderModules = Object.fromEntries(
      Object.entries({
        orangeFragWGSL,
        texturedFragWGSL,
        distanceSizedPointsVertWGSL,
        fixedSizePointsVertWGSL,
      }).map(([key, code]) => [key, device!.createShaderModule({ code })])
    );

    const fragModules = [
      shaderModules.orangeFragWGSL,
      shaderModules.texturedFragWGSL,
    ];

    const vertModules = [
      shaderModules.distanceSizedPointsVertWGSL,
      shaderModules.fixedSizePointsVertWGSL,
    ];

    const depthFormat = "depth24plus";

    // make pipelines for each combination
    const pipelines = vertModules.map((vertModule) =>
      fragModules.map((fragModule) =>
        device!.createRenderPipeline({
          layout: pipelineLayout,
          vertex: {
            module: vertModule,
            buffers: [
              {
                arrayStride: 3 * 4, // 3 floats, 4 bytes each
                stepMode: "instance",
                attributes: [
                  { shaderLocation: 0, offset: 0, format: "float32x3" }, // position
                ],
              },
            ],
          },
          fragment: {
            module: fragModule,
            targets: [
              {
                format: presentationFormat,
                blend: {
                  color: {
                    srcFactor: "one",
                    dstFactor: "one-minus-src-alpha",
                  },
                  alpha: {
                    srcFactor: "one",
                    dstFactor: "one-minus-src-alpha",
                  },
                },
              },
            ],
          },
          depthStencil: {
            depthWriteEnabled: true,
            depthCompare: "less",
            format: depthFormat,
          },
        })
      )
    );

    const vertexData = createFibonacciSphereVertices({
      radius: 1,
      numSamples: 1000,
    });
    const kNumPoints = vertexData.length / 3;

    const vertexBuffer = device!.createBuffer({
      label: "vertex buffer vertices",
      size: vertexData.byteLength,
      usage: GPUBufferUsage.VERTEX | GPUBufferUsage.COPY_DST,
    });
    device!.queue.writeBuffer(vertexBuffer, 0, vertexData);

    const uniformValues = new Float32Array(16 + 2 + 1 + 1);
    const uniformBuffer = device!.createBuffer({
      size: uniformValues.byteLength,
      usage: GPUBufferUsage.UNIFORM | GPUBufferUsage.COPY_DST,
    });
    const kMatrixOffset = 0;
    const kResolutionOffset = 16;
    const kSizeOffset = 18;
    const matrixValue = uniformValues.subarray(
      kMatrixOffset,
      kMatrixOffset + 16
    );
    const resolutionValue = uniformValues.subarray(
      kResolutionOffset,
      kResolutionOffset + 2
    );
    const sizeValue = uniformValues.subarray(kSizeOffset, kSizeOffset + 1);

    // Use canvas 2d to make texture data
    const ctx = new OffscreenCanvas(64, 64).getContext("2d");
    ctx!.font = "60px sans-serif";
    ctx!.textAlign = "center";
    ctx!.textBaseline = "middle";
    ctx!.fillText("🦋", 32, 32);

    const sampler = device!.createSampler();
    const texture = device!.createTexture({
      size: [ctx!.canvas.width, ctx!.canvas.height],
      format: "rgba8unorm",
      usage:
        GPUTextureUsage.COPY_DST |
        GPUTextureUsage.TEXTURE_BINDING |
        GPUTextureUsage.RENDER_ATTACHMENT,
    });
    device!.queue.copyExternalImageToTexture(
      { source: ctx!.canvas, flipY: true },
      { texture },
      [ctx!.canvas.width, ctx!.canvas.height]
    );

    const bindGroup = device!.createBindGroup({
      layout: bindGroupLayout,
      entries: [
        { binding: 0, resource: { buffer: uniformBuffer } },
        { binding: 1, resource: sampler },
        { binding: 2, resource: texture.createView() },
      ],
    });

    const renderPassDescriptor: GPURenderPassDescriptor = {
      label: "our basic canvas renderPass",
      colorAttachments: [
        {
          view: undefined as unknown as GPUTextureView, // assigned later
          clearValue: [0.3, 0.3, 0.3, 1],
          loadOp: "clear",
          storeOp: "store",
        },
      ],
      depthStencilAttachment: {
        view: undefined as unknown as GPUTextureView, // to be filled out when we render
        depthClearValue: 1.0,
        depthLoadOp: "clear",
        depthStoreOp: "store",
      },
    };

    // const gui = new GUI();
    // gui.add(settings, "fixedSize");
    // gui.add(settings, "textured");
    // gui.add(settings, "size", 0, 80);

    let depthTexture: GPUTexture;

    function render(time: number) {
      // Convert to seconds.
      time *= 0.001;

      // Get the current texture from the canvas context and
      // set it as the texture to render to.
      const canvasTexture = context!.getCurrentTexture();
      // @ts-ignore
      renderPassDescriptor.colorAttachments[0].view =
        canvasTexture.createView();

      // If we don't have a depth texture OR if its size is different
      // from the canvasTexture when make a new depth texture
      if (
        !depthTexture ||
        depthTexture.width !== canvasTexture.width ||
        depthTexture.height !== canvasTexture.height
      ) {
        if (depthTexture) {
          depthTexture.destroy();
        }
        depthTexture = device!.createTexture({
          size: [canvasTexture.width, canvasTexture.height],
          format: depthFormat,
          usage: GPUTextureUsage.RENDER_ATTACHMENT,
        });
      }
      // @ts-ignore
      renderPassDescriptor.depthStencilAttachment.view =
        depthTexture.createView();

      const { size, fixedSize, textured } = self.settings;

      const pipeline = pipelines[fixedSize ? 1 : 0][textured ? 1 : 0];

      // Set the size in the uniform values
      sizeValue[0] = size;

      const fov = (90 * Math.PI) / 180;
      const aspect = canvas!.clientWidth / canvas!.clientHeight;
      const projection = mat4.perspective(fov, aspect, 0.1, 50);
      const view = mat4.lookAt(
        [0, 0, 1.5], // position
        [0, 0, 0], // target
        [0, 1, 0] // up
      );
      const viewProjection = mat4.multiply(projection, view);
      mat4.rotateY(viewProjection, time, matrixValue);
      mat4.rotateX(matrixValue, time * 0.1, matrixValue);

      // Update the resolution in the uniform values
      resolutionValue.set([canvasTexture.width, canvasTexture.height]);

      // Copy the uniform values to the GPU
      device!.queue.writeBuffer(uniformBuffer, 0, uniformValues);

      const encoder = device!.createCommandEncoder();
      const pass = encoder.beginRenderPass(renderPassDescriptor);
      pass.setPipeline(pipeline);
      pass.setVertexBuffer(0, vertexBuffer);
      pass.setBindGroup(0, bindGroup);
      pass.draw(6, kNumPoints);
      pass.end();

      const commandBuffer = encoder.finish();
      device!.queue.submit([commandBuffer]);

      requestAnimationFrame(render);
    }

    requestAnimationFrame(render);
  }

  onUpdate(data: any): void {
      this.settings = data;
  }
}
